import torch
import os
import utils
from model.hidden import Hidden
this_run_folder = "/home/zzc/Experience/HiDDeN_zzc/runs/run_Jpeg-50-no-normalize_2024-03-08_08-32-37-100"
device = torch.device('cuda:3') if torch.cuda.is_available() else torch.device('cpu')
options_file = os.path.join(this_run_folder, 'options-and-config.pickle')
train_options, hidden_config, noise_config = utils.load_options(options_file)

last_checkpoint_file = "/home/zzc/Experience/HiDDeN_zzc/runs/run_Jpeg-50-no-normalize_2024-03-08_08-32-37-100/checkpoints/run_Jpeg-50-no-normalize_2024-03-08_08-32-37-100--epoch-700.pyt"
checkpoint = torch.load(last_checkpoint_file)
train_options.start_epoch = checkpoint['epoch'] + 1

# 加载模型参数
model = Hidden(hidden_config, device, tb_logger=None)
model.encoder_decoder.load_state_dict(checkpoint['enc-dec-model'])
model.encoder_decoder.eval()
# 将参数信息保存到字符串中
param_info = ""
for name, param in model.encoder_decoder.named_parameters():
    param_info += f"Parameter name: {name}, Shape: {param.shape}\n"
    param_info += f"Value: {param}\n"

# 将参数信息写入到txt文件中
with open('model_params.txt', 'w') as f:
    f.write(param_info)

print("Model parameters info saved to 'model_params.txt'")